'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: pspnet.py
@time: 2020-06-08 14:19:50
@desc: 
'''
from ELib.model.model_zoo import MODEL_ZOO
from ELib.model.segbase import SegBaseModel
from ELib.basic_model import _FCNHead, PyramidPooling
import torch


@MODEL_ZOO.register()
class PSPNet(SegBaseModel):
    r"""Pyramid Scene Parsing Network
    Reference:
        Zhao, Hengshuang, Jianping Shi, Xiaojuan Qi, Xiaogang Wang, and Jiaya Jia.
        "Pyramid scene parsing network." *CVPR*, 2017
    """

    def __init__(self, cfg):
        super(PSPNet, self).__init__(cfg)
        self.head = _PSPHead(self.nclass)
        if self.aux:
            self.auxlayer = _FCNHead(1024, self.nclass)

        self.__setattr__('decoder', ['head', 'auxlayer'] if self.aux else ['head'])

    def forward(self, x):
        size = x.size()[2:]
        _, _, c3, c4 = self.encoder(x)
        outputs = []
        x = self.head(c4)
        x = torch.nn.functional.interpolate(x, size, mode='bilinear', align_corners=True)
        outputs.append(x)

        if self.aux:
            auxout = self.auxlayer(c3)
            auxout = torch.nn.functional.interpolate(auxout, size, mode='bilinear', align_corners=True)
            outputs.append(auxout)
        return tuple(outputs)


class _PSPHead(torch.nn.Module):
    def __init__(self, nclass, norm_layer=torch.nn.BatchNorm2d, norm_kwargs=None, **kwargs):
        super(_PSPHead, self).__init__()
        self.psp = PyramidPooling(2048, norm_layer=norm_layer)
        self.block = torch.nn.Sequential(
            torch.nn.Conv2d(4096, 512, 3, padding=1, bias=False),
            norm_layer(512, **({} if norm_kwargs is None else norm_kwargs)),
            torch.nn.ReLU(True),
            torch.nn.Dropout(0.1),
            torch.nn.Conv2d(512, nclass, 1)
        )

    def forward(self, x):
        x = self.psp(x)
        return self.block(x)